# Copyright (c) 2022-2024, The Isaac Lab Project Developers.
# All rights reserved.
#
# SPDX-License-Identifier: BSD-3-Clause

from __future__ import annotations

import gymnasium as gym
import torch

import isaaclab.sim as sim_utils
from isaaclab.assets import Articulation, ArticulationCfg
from isaaclab.envs import DirectRLEnv, DirectRLEnvCfg
from isaaclab.envs.ui import BaseEnvWindow
from isaaclab.markers import VisualizationMarkers
from isaaclab.scene import InteractiveSceneCfg
from isaaclab.sim import SimulationCfg
from isaaclab.terrains import TerrainImporterCfg
from isaaclab.utils import configclass
from isaaclab.utils.math import subtract_frame_transforms, random_yaw_orientation, compute_pose_error, yaw_angle
from isaaclab.sensors import TiledCamera, TiledCameraCfg, ContactSensor, ContactSensorCfg
from isaaclab.sensors.ray_caster import RayCaster, RayCasterCfg, patterns, RTXRayCasterCfg, RTXRayCaster
from isaaclab.sensors import RangeSensor, RangeSensorCfg
from isaaclab.assets import (
    Articulation,
    ArticulationCfg,
    AssetBaseCfg,
    RigidObject,
    RigidObjectCfg,
    RigidObjectCollection,
    RigidObjectCollectionCfg,
)
##
# Pre-defined configs
##
from isaaclab_assets import UAV_LIDAR_CFG  # isort: skip
from isaaclab.markers import CUBOID_MARKER_CFG  # isort: skip
from isaaclab.markers.config import BLUE_ARROW_X_MARKER_CFG, GREEN_ARROW_X_MARKER_CFG, RED_ARROW_X_MARKER_CFG

from .gen_random_path import Gen3DTrajectories

class UAVEnvWindow(BaseEnvWindow):
    """Window manager for the UAV environment."""

    def __init__(self, env: UAVTrackingEnv, window_name: str = "IsaacLab"):
        """Initialize the window.

        Args:
            env: The environment object.
            window_name: The name of the window. Defaults to "IsaacLab".
        """
        # initialize base window
        super().__init__(env, window_name)
        # add custom UI elements
        with self.ui_window_elements["main_vstack"]:
            with self.ui_window_elements["debug_frame"]:
                with self.ui_window_elements["debug_vstack"]:
                    # add command manager visualization
                    self._create_debug_vis_ui_element("targets", self.env)


@configclass
class UAVTrackingEnvCfg(DirectRLEnvCfg):
    # env
    episode_length_s = 200.0
    decimation = 2
    action_space = 4
    state_space = 0
    debug_vis = True

    ui_window_class_type = UAVEnvWindow

    # simulation
    sim: SimulationCfg = SimulationCfg(
        dt=1 / 100,
        render_interval=decimation,
        enable_scene_query_support = True,
        use_fabric=True, 
        physics_material=sim_utils.RigidBodyMaterialCfg(
            friction_combine_mode="multiply",
            restitution_combine_mode="multiply",
            static_friction=1.0,
            dynamic_friction=1.0,
            restitution=0.0,
        ),
    )
    terrain = TerrainImporterCfg(
        prim_path="/World/ground",
        terrain_type="plane",
        collision_group=-1,
        physics_material=sim_utils.RigidBodyMaterialCfg(
            friction_combine_mode="multiply",
            restitution_combine_mode="multiply",
            static_friction=1.0,
            dynamic_friction=1.0,
            restitution=0.0,
        ),
        debug_vis=False,
    )

    # scene
    scene: InteractiveSceneCfg = InteractiveSceneCfg(num_envs=4096, env_spacing=25.0, replicate_physics=True)

    physx_lidar = RangeSensorCfg(
        prim_path="/World/envs/env_.*/Robot/body/Lidar",
        # update_period=0.025,  # Update rate of 40Hz
        # data_types=["point_cloud"],  # Assuming the LiDAR generates point cloud data
        horizontal_fov=360.0,  # Horizontal field of view of 270 degrees
        horizontal_resolution=0.4,  # Horizontal resolution of 0.5 degrees
        max_range=30.0,  # Maximum range of 30 meters
        min_range=0.020,  # Minimum range of 0.1 meters
        rotation_rate=0.0,  # Rotation rate of 0.0 radians per second
        offset=RangeSensorCfg.OffsetCfg(
            pos=(0.0, 0.12, 0.0),  # Example position offset from the robot base
            rot=(0.7071, 0.7071, 0.0, 0.0),  # Example rotation offset; no rotation in this case
            convention="ros"  # Frame convention
        ),
        draw_lines=False,
        draw_points=True,
    )

    # SICK_TIM781
    # info: RTXRayCasterInfo(numChannels=811, numEchos=1, numReturnsPerScan=811, renderProductPath='/Render/OmniverseKit/HydraTextures/Replicator_01', ticksPerScan=1)

    # object collection
    obstacle_cfg: RigidObjectCollectionCfg = RigidObjectCollectionCfg(
        rigid_objects={
            "obs_A": RigidObjectCfg(
                prim_path="/World/envs/env_.*/Obs_A",
                spawn=sim_utils.CuboidCfg(
                    size=(0.5, 0.5, 20.0),
                    visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.0, 0.0, 1.0), metallic=0.2),
                    rigid_props=sim_utils.RigidBodyPropertiesCfg(
                        solver_position_iteration_count=4, solver_velocity_iteration_count=0
                    ),
                    mass_props=sim_utils.MassPropertiesCfg(mass=10000.0),
                    collision_props=sim_utils.CollisionPropertiesCfg(),
                ),
                init_state=RigidObjectCfg.InitialStateCfg(pos=(5.0, 5.0, 10.0)),
            ),
            "obs_B": RigidObjectCfg(
                prim_path="/World/envs/env_.*/Obs_B",
                spawn=sim_utils.CuboidCfg(
                    size=(0.5, 0.7, 20.0),
                    visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(1.0, 1.0, 0.0), metallic=0.2),
                    rigid_props=sim_utils.RigidBodyPropertiesCfg(
                        solver_position_iteration_count=4, solver_velocity_iteration_count=0
                    ),
                    mass_props=sim_utils.MassPropertiesCfg(mass=10000.0),
                    collision_props=sim_utils.CollisionPropertiesCfg(),
                ),
                init_state=RigidObjectCfg.InitialStateCfg(pos=(5.0, 0.0, 10.0)),
            ),
            "obs_C": RigidObjectCfg(
                prim_path="/World/envs/env_.*/Obs_C",
                spawn=sim_utils.CylinderCfg(
                    radius=0.35,
                    height=20.0,
                    visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.0, 0.0, 1.0), metallic=0.2),
                    rigid_props=sim_utils.RigidBodyPropertiesCfg(
                        solver_position_iteration_count=4, solver_velocity_iteration_count=0
                    ),
                    mass_props=sim_utils.MassPropertiesCfg(mass=10000.0),
                    collision_props=sim_utils.CollisionPropertiesCfg(),
                ),
                init_state=RigidObjectCfg.InitialStateCfg(pos=(5.0, -5.0, 10.0)),
            ),
            "obs_D": RigidObjectCfg(
                prim_path="/World/envs/env_.*/Obs_D",
                spawn=sim_utils.CylinderCfg(
                    radius=0.25,
                    height=20.0,
                    visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(1.0, 0.0, 1.0), metallic=0.2),
                    rigid_props=sim_utils.RigidBodyPropertiesCfg(
                        solver_position_iteration_count=4, solver_velocity_iteration_count=0
                    ),
                    mass_props=sim_utils.MassPropertiesCfg(mass=10000.0),
                    collision_props=sim_utils.CollisionPropertiesCfg(),
                ),
                init_state=RigidObjectCfg.InitialStateCfg(pos=(-5.0, -5.0, 10.0)),
            ),
            "obs_E": RigidObjectCfg(
                prim_path="/World/envs/env_.*/Obs_E",
                spawn=sim_utils.CylinderCfg(
                    radius=0.35,
                    height=20.0,
                    visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(1.0, 1.0, 0.0), metallic=0.2),
                    rigid_props=sim_utils.RigidBodyPropertiesCfg(
                        solver_position_iteration_count=4, solver_velocity_iteration_count=0
                    ),
                    mass_props=sim_utils.MassPropertiesCfg(mass=10000.0),
                    collision_props=sim_utils.CollisionPropertiesCfg(),
                ),
                init_state=RigidObjectCfg.InitialStateCfg(pos=(-5.0, 0.0, 10.0)),
            ),
            "obs_F": RigidObjectCfg(
                prim_path="/World/envs/env_.*/Obs_F",
                spawn=sim_utils.CuboidCfg(
                    size=(0.7, 0.5, 20.0),
                    visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.0, 1.0, 0.0), metallic=0.2),
                    rigid_props=sim_utils.RigidBodyPropertiesCfg(
                        solver_position_iteration_count=4, solver_velocity_iteration_count=0
                    ),
                    mass_props=sim_utils.MassPropertiesCfg(mass=10000.0),
                    collision_props=sim_utils.CollisionPropertiesCfg(),
                ),
                init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, -5.0, 10.0)),
            ),
            "obs_G": RigidObjectCfg(
                prim_path="/World/envs/env_.*/Obs_G",
                spawn=sim_utils.CuboidCfg(
                    size=(0.5, 0.5, 20.0),
                    visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.0, 1.0, 1.0), metallic=0.2),
                    rigid_props=sim_utils.RigidBodyPropertiesCfg(
                        solver_position_iteration_count=4, solver_velocity_iteration_count=0
                    ),
                    mass_props=sim_utils.MassPropertiesCfg(mass=10000.0),
                    collision_props=sim_utils.CollisionPropertiesCfg(),
                ),
                init_state=RigidObjectCfg.InitialStateCfg(pos=(-5.0, 5.0, 10.0)),
            ),
            "obs_H": RigidObjectCfg(
                prim_path="/World/envs/env_.*/Obs_H",
                spawn=sim_utils.CuboidCfg(
                    size=(0.5, 0.5, 20.0),
                    visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(1, 1.0, 1.0), metallic=0.2),
                    rigid_props=sim_utils.RigidBodyPropertiesCfg(
                        solver_position_iteration_count=4, solver_velocity_iteration_count=0
                    ),
                    mass_props=sim_utils.MassPropertiesCfg(mass=10000.0),
                    collision_props=sim_utils.CollisionPropertiesCfg(),
                ),
                init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 5.0, 10.0)),
            ),
        }
    )

    # robot
    robot: ArticulationCfg = UAV_LIDAR_CFG.replace(prim_path="/World/envs/env_.*/Robot")
    thrust_to_weight = 10.0
    moment_scale = 0.05
    distance_threshold = 10.0
    # 设置碰撞距离阈值
    collision_threshold = 0.5

    # reward scales
    lin_vel_reward_scale = -0.05
    ang_vel_reward_scale = -0.05
    distance_to_goal_reward_scale = 15.0
    collision_reward_scale = -1.0
    action_rate_reward_scale = -0.1
    goal_reached_reward_scale = 0.5
    delta_goal_reward_scale = 200.0
    angle_reward_scale = 4.0
    target_angle_reward_scale = 2.0

    observation_space = {
        "robot-state": 20,
        "lidar": 900,
    }

class UAVTrackingEnv(DirectRLEnv):
    cfg: UAVTrackingEnvCfg

    def __init__(self, cfg: UAVTrackingEnvCfg, render_mode: str | None = None, **kwargs):
        super().__init__(cfg, render_mode, **kwargs)

        # Total thrust and moment applied to the base of the uav
        self._actions = torch.zeros(self.num_envs, gym.spaces.flatdim(self.single_action_space), device=self.device)
        self._previous_actions = torch.zeros(self.num_envs, gym.spaces.flatdim(self.single_action_space), device=self.device)

        self._thrust = torch.zeros(self.num_envs, 1, 3, device=self.device)
        self._moment = torch.zeros(self.num_envs, 1, 3, device=self.device)
        # Goal position
        self._desired_pos_w = torch.zeros(self.num_envs, 3, device=self.device)
        self._collision = torch.zeros(self.num_envs, device=self.device)
        self._is_goal_reached = torch.zeros(self.num_envs, device=self.device, dtype=torch.bool)
        self._dist_to_goal = torch.zeros(self.num_envs, device=self.device)
        self._previous_dist_to_goal = torch.zeros(self.num_envs, device=self.device)
        self._dist_to_obstacles = torch.zeros(self.num_envs, device=self.device)
        
        # Logging
        self._episode_sums = {
            key: torch.zeros(self.num_envs, dtype=torch.float, device=self.device)
            for key in [
                "lin_vel",
                "ang_vel",
                "distance_to_goal",
                "collision",
                "action_rate_l2",
                "goal_reached",
                "delta_goal_dist",
                "safety_static",
                "reverse_penalty",
                "angle_error",
            ]
        }
        # Get specific body indices
        self._body_id = self._robot.find_bodies("body")[0]
        self._robot_mass = self._robot.root_physx_view.get_masses()[0].sum()
        self._gravity_magnitude = torch.tensor(self.sim.cfg.gravity, device=self.device).norm()
        self._robot_weight = (self._robot_mass * self._gravity_magnitude).item()

        self.trajectories_controller = Gen3DTrajectories(num_objects=self.num_envs, device=self.device)
        # add handle for debug visualization (this is set to a valid handle inside set_debug_vis)
        self.set_debug_vis(self.cfg.debug_vis)

    def _setup_scene(self):
        self._robot = Articulation(self.cfg.robot)
        self.scene.articulations["robot"] = self._robot

        # add physX lidar
        self._physx_lidar = RangeSensor(self.cfg.physx_lidar)
        self.scene.sensors["physx_lidar"] = self._physx_lidar

        # obstacles
        self._obstacles = RigidObjectCollection(self.cfg.obstacle_cfg)
        self.scene.rigid_object_collections["obstacles"] = self._obstacles

        self.cfg.terrain.num_envs = self.scene.cfg.num_envs
        self.cfg.terrain.env_spacing = self.scene.cfg.env_spacing
        self._terrain = self.cfg.terrain.class_type(self.cfg.terrain)
        # clone, filter, and replicate
        self.scene.clone_environments(copy_from_source=False)
        self.scene.filter_collisions(global_prim_paths=[self.cfg.terrain.prim_path])
        # add lights
        light_cfg = sim_utils.DomeLightCfg(intensity=2000.0, color=(0.75, 0.75, 0.75))
        light_cfg.func("/World/Light", light_cfg)

    def _pre_physics_step(self, actions: torch.Tensor):
        self._actions = actions.clone().clamp(-1.0, 1.0)
        self._thrust[:, 0, 1] = self.cfg.thrust_to_weight * self._robot_weight * (self._actions[:, 0] + 1.0) / 2.0
        self._moment[:, 0, :] = self.cfg.moment_scale * self._actions[:, 1:]

    def _apply_action(self):
        self._robot.set_external_force_and_torque(self._thrust, self._moment, body_ids=self._body_id)

    def _get_observations(self) -> dict:
        self._previous_actions = self._actions.clone()
        self.update_desired_pos()
        pos_error, axis_angle_error = compute_pose_error(self._desired_pos_w, self._desired_quat_w, self._robot.data.root_pos_w, self._robot.data.root_quat_w)
        yaw_robo = yaw_angle(self._robot.data.root_quat_w)
        yaw_desired = yaw_angle(self._desired_quat_w)

        robot_state = torch.cat(
            [
                self._robot.data.root_lin_vel_b,
                self._robot.data.root_ang_vel_b,
                self._robot.data.projected_gravity_b,
                pos_error,
                yaw_robo.unsqueeze(1),
                yaw_desired.unsqueeze(1),
                axis_angle_error[:, -1:],
                self._dist_to_goal.unsqueeze(1),
                self._actions,
            ],
            dim=-1,
        )

        observations = {
            "policy": {
                "robot-state": robot_state,
                "lidar": self._physx_lidar.data.output["linear_depth"],
            }
        }

        # observations = {"policy": obs}
        return observations

    def _get_rewards(self) -> torch.Tensor:
        lin_vel = torch.sum(torch.square(self._robot.data.root_lin_vel_b), dim=1)
        ang_vel = torch.sum(torch.square(self._robot.data.root_ang_vel_b), dim=1)
        self._dist_to_goal = torch.linalg.norm(self._desired_pos_w - self._robot.data.root_pos_w, dim=1)
        distance_to_goal_mapped = 1 - torch.tanh(0.5 * self._dist_to_goal)
        delta_goal_dist = self._previous_dist_to_goal - self._dist_to_goal
        # print("delta_goal_dist:", delta_goal_dist)
        self._previous_dist_to_goal = self._dist_to_goal.clone()
        # print("_dist_to_goal:", self._dist_to_goal)

        # action rate
        action_rate = torch.sum(torch.square(self._actions - self._previous_actions), dim=1)
        self._is_goal_reached = self._dist_to_goal < 0.05
        yaw_robo = yaw_angle(self._robot.data.root_quat_w)
        # yaw_desired = yaw_angle(self._desired_quat_w)
        # angle_error = wrap_to_pi(yaw_robo - yaw_desired)
        pos_error, axis_angle_error = compute_pose_error(self._desired_pos_w, self._desired_quat_w, self._robot.data.root_pos_w, self._robot.data.root_quat_w)
        # print("pos_error:", pos_error, "axis_angle_error:", axis_angle_error, axis_angle_error[:,2], "angle_error_mapped:", torch.cos(axis_angle_error[:,2]))
        target_angle = torch.atan2(pos_error[:, 1], pos_error[:, 0])
        # print("target_angle:", target_angle)
        angel_error_to_target = yaw_robo - target_angle + torch.pi
        # print("distance_to_goal_mapped:", distance_to_goal_mapped)
        # print("angel_error_to_target:", angel_error_to_target, torch.cos(angel_error_to_target))
        angle_error_mapped = torch.cos(axis_angle_error[:,2])
        # target_angle_reward = torch.cos(angel_error_to_target) * torch.tanh(self._dist_to_goal)
        # print("target_angle_reward:",  target_angle_reward)
        # reverse_penalty = (self._robot.data.root_lin_vel_b[:, 0] < 0).float() * (-1.0)

        lidar_data = self._physx_lidar.data.output["linear_depth"]
        min_dist, _ = torch.min(lidar_data, dim=1)
        # print("min_dist:", min_dist)

        self._dist_to_obstacles = min_dist
        self._collision = self._dist_to_obstacles < self.cfg.collision_threshold
        # print("self._collision:", self._collision)

        # 计算剩余空间（从障碍物到LiDAR最大探测距离）
        remaining_space = 1.65 - min_dist
        # 对剩余空间取自然对数，并计算平均值作为奖励
        reward_safety_static = torch.pow(torch.tensor(10.0), remaining_space.clamp(min=1e-6, max=1.65))-1

        rewards = {
            "lin_vel": lin_vel * self.cfg.lin_vel_reward_scale * self.step_dt,
            "ang_vel": ang_vel * self.cfg.ang_vel_reward_scale * self.step_dt,
            "distance_to_goal": distance_to_goal_mapped * self.cfg.distance_to_goal_reward_scale * self.step_dt,
            "action_rate_l2": action_rate * self.cfg.action_rate_reward_scale * self.step_dt,
            "goal_reached": self._is_goal_reached.float() * self.cfg.goal_reached_reward_scale,
            "delta_goal_dist": delta_goal_dist * self.cfg.delta_goal_reward_scale * self.step_dt,
            "angle_error":  angle_error_mapped * distance_to_goal_mapped* self.cfg.angle_reward_scale * self.step_dt,
            "collision": self._collision.float() * self.cfg.collision_reward_scale,
            "safety_static" : reward_safety_static * self.step_dt,
            # "reverse_penalty": reverse_penalty * self.step_dt,
            # "target_angle": target_angle_reward * self.cfg.target_angle_reward_scale * self.step_dt,
        }
        reward = torch.sum(torch.stack(list(rewards.values())), dim=0)
        # Logging
        for key, value in rewards.items():
            self._episode_sums[key] += value
        return reward

    def _get_dones(self) -> tuple[torch.Tensor, torch.Tensor]:
        time_out = self.episode_length_buf >= self.max_episode_length - 1
        died = torch.logical_or(self._dist_to_goal > self.cfg.distance_threshold, self._collision)
        # died = time_out
        return died, time_out

    def _reset_idx(self, env_ids: torch.Tensor | None):
        if env_ids is None or len(env_ids) == self.num_envs:
            env_ids = self._robot._ALL_INDICES

        # Logging
        final_distance_to_goal = torch.linalg.norm(
            self._desired_pos_w[env_ids] - self._robot.data.root_pos_w[env_ids], dim=1
        ).mean()
        extras = dict()
        for key in self._episode_sums.keys():
            episodic_sum_avg = torch.mean(self._episode_sums[key][env_ids])
            extras["Episode_Reward/" + key] = episodic_sum_avg / self.max_episode_length_s
            self._episode_sums[key][env_ids] = 0.0
        self.extras["log"] = dict()
        self.extras["log"].update(extras)
        extras = dict()
        extras["Episode_Termination/died"] = torch.count_nonzero(self.reset_terminated[env_ids]).item()
        extras["Episode_Termination/time_out"] = torch.count_nonzero(self.reset_time_outs[env_ids]).item()
        extras["Metrics/final_distance_to_goal"] = final_distance_to_goal.item()
        self.extras["log"].update(extras)

        self._robot.reset(env_ids)
        super()._reset_idx(env_ids)
        self._collision[env_ids] = 0
        if len(env_ids) == self.num_envs:
            # Spread out the resets to avoid spikes in training when many environments reset at a similar time
            self.episode_length_buf = torch.randint_like(self.episode_length_buf, high=int(self.max_episode_length))

        self._actions[env_ids] = 0.0
        self._previous_actions[env_ids] = 0.0
        self._is_goal_reached[env_ids] &= False
        # Sample new commands
        self.update_desired_pos()
        # Reset robot state
        joint_pos = self._robot.data.default_joint_pos[env_ids]
        joint_vel = self._robot.data.default_joint_vel[env_ids]
        default_root_state = self._robot.data.default_root_state[env_ids]
        default_root_state[:, :3] = self._desired_pos_w[env_ids]
        default_root_state[:, 3:7] = self._desired_quat_w[env_ids]
        self._robot.write_root_pose_to_sim(default_root_state[:, :7], env_ids)
        self._robot.write_root_velocity_to_sim(default_root_state[:, 7:], env_ids)
        self._robot.write_joint_state_to_sim(joint_pos, joint_vel, None, env_ids)
        self._previous_dist_to_goal[env_ids] = torch.linalg.norm(self._desired_pos_w[env_ids] - self._robot.data.root_pos_w[env_ids], dim=1)

    def _set_debug_vis_impl(self, debug_vis: bool):
        # set visibility of markers
        # note: parent only deals with callbacks. not their visibility
        if debug_vis:
            # create markers if necessary for the first tome
            if not hasattr(self, "goal_pos_visualizer"):
                # -- goal
                marker_cfg = GREEN_ARROW_X_MARKER_CFG.copy()
                marker_cfg.prim_path = "/Visuals/Command/goal_position"
                marker_cfg.markers["arrow"].scale = (0.1, 0.1, 0.3)
                self.base_target_goal_visualizer = VisualizationMarkers(marker_cfg)
                # -- current
                marker_cfg = BLUE_ARROW_X_MARKER_CFG.copy()
                marker_cfg.prim_path = "/Visuals/Command/robot_current"
                marker_cfg.markers["arrow"].scale = (0.04, 0.04, 0.14)
                self.base_robot_visualizer = VisualizationMarkers(marker_cfg)
            # set their visibility to true
            self.base_target_goal_visualizer.set_visibility(True)
            self.base_robot_visualizer.set_visibility(True)
        else:
            if hasattr(self, "goal_pos_visualizer"):
                self.base_target_goal_visualizer.set_visibility(False)
                self.base_robot_visualizer.set_visibility(False)

    def _debug_vis_callback(self, event):
        # get marker location
        # -- base state
        robot_pos_w = self._robot.data.root_pos_w.clone()
        robot_quat_w = self._robot.data.root_quat_w.clone()
        robot_pos_w[:, 2] += 0.1
        # -- resolve the scales and quaternions
        # vel_des_arrow_scale, vel_des_arrow_quat = self._resolve_xy_velocity_to_arrow(self.command[:, :2])
        # vel_arrow_scale, vel_arrow_quat = self._resolve_xy_velocity_to_arrow(self.robot.data.root_lin_vel_b[:, :2])
        # display markers
        self.base_target_goal_visualizer.visualize(self._desired_pos_w, self._desired_quat_w)
        self.base_robot_visualizer.visualize(robot_pos_w, robot_quat_w)

    def update_desired_pos(self):
        self._desired_pos_w, self._desired_quat_w = self.trajectories_controller.update()